#!/usr/bin/env python

from __future__ import with_statement
import sys
import plistlib
import re

# Helper functions.
def find_node(plist, userinfo_key, userinfo_value):
    for x in plist['GraphicsList']:
        if x.has_key('UserInfo') and x['UserInfo'].has_key(userinfo_key) and x['UserInfo'][userinfo_key] == userinfo_value:
            return x
    return None

def set_color(plist, userinfo_key, taskid, color):
    node = find_node(plist, userinfo_key, taskid)
    if node == None:
        print 'Unable to find node "%s" in graph.'%taskid
        sys.exit(1)
    node['Style']['fill']['Color']['r'] = color[0]
    node['Style']['fill']['Color']['g'] = color[1]
    node['Style']['fill']['Color']['b'] = color[2]

colors = [('0','0','1'), 
          ('0','1','1'), 
          ('0','1','0'),
          ('1','1','0'),
          ('1','0','0'),
          ('1','0','1'), 
          ('0.298039','0.298039','0.298039'), 
          ('0.6','0.6','0.6')]
color_mapping = {}
def find_color(peer):
    global colors, color_mapping
    if not color_mapping.has_key(peer):
        color_mapping[peer] = colors.pop()
    return color_mapping[peer]

# Check command line arguments.
if len(sys.argv) < 4:
    print "Usage: colorize_graph.py graffle-input logfile graffle-output"
    sys.exit(1)

# Read in the graffle file.
graffle_data = None
with open(sys.argv[1], 'r') as infile:
    graffle_data = plistlib.readPlist(infile)

# Read in the log file.
log_lines = None
with open(sys.argv[2], 'r') as infile:
    log_lines = infile.readlines()

# Parse the log.
regexp = re.compile("([a-z]+)(\d+)")
x = 1
y = 24
while y <= len(log_lines):
    for line in log_lines[y-24:y]:
        line = line.strip()
        peer, taskid = line.split(' -> ')
        m = regexp.match(taskid)
        taskid = m.group(1) + str(int(m.group(2))%8)
        # Set node color.
        color = find_color(peer)
        set_color(graffle_data, 'taskid', taskid, color)
        # Set text on node.
        node = find_node(graffle_data, 'taskid', taskid)
        text = node['Text']['Text']
        text = text.replace('foo', '%i'%x)
        node['Text']['Text'] = text
        x += 1

    # Write the legend.
    peers = color_mapping.keys()
    peers.sort()
    x = 1
    for peer in peers:
        color = color_mapping[peer]
        set_color(graffle_data, 'label', '%i'%x, color)
        text_label = find_node(graffle_data, 'tlabel', '%i'%x)
        text = text_label['Text']['Text']
        text = text.replace('label%i'%x, peer)
        text_label['Text']['Text'] = text
        x += 1

    # Write the output graffle.
    plistlib.writePlist(graffle_data, sys.argv[3] + str(y/24) + '.graffle')

    y += 24

